from nets.yolo3 import yolo_body
from keras.layers import Input
from yolo import YOLO
from PIL import Image
from cv import videoToImg
from cv import mkdir
import cv2
import os.path
import os
import json

def print_split():
    print("*************************************")


def cropImg(img, box):                                    #按照尺寸裁剪图片
    cropped = img.crop((box[1], box[0], box[3], box[2]))
    return cropped

def getCropImgPath(video, box_class):
    dir_name = 'D:/yolo3-keras/croppedImg/' + video
    if not os.path.exists(dir_name):
        os.mkdir(dir_name)
    _count = traffic_count[box_class]
    crapPath = dir_name + '/' + class_names[box_class] + str(_count) + '.jpg'
    return crapPath


# 检测车流量
def traffic(old_boxs, new_boxs, old_classes, new_classes, img, frame):
    for i in range(len(new_boxs)):
        current_box = new_boxs[i]                            #当前box
        current_x = (current_box[1] + current_box[3]) / 2    #求x中点
        current_y = (current_box[0] + current_box[2]) / 2    #求y中点
        current_width = current_box[3] - current_box[1]      #求宽度
        current_height = current_box[2] - current_box[0]     #求高度
        current_class = new_classes[i]                       #当前种类
        for j in range(len(old_boxs)):
            old_box = old_boxs[j]
            old_x = (old_box[1] + old_box[3]) / 2
            old_y = (old_box[0] + old_box[2]) / 2
            old_class = old_classes[j]
            # 判断是否在判定范围内，并且为相同class
            if ((current_class == old_class) and (abs(current_x - old_x) < current_width) and (abs(current_y - old_y) < current_height)):
                # 判断是否刚超过baseline
                if (((current_x - img_width * baseline_x) * (old_x - img_width * baseline_x) < 0)):
                    if (current_y > img_width * 0.7):
                        continue
                    traffic_count[current_class] += 1
                    cropped = cropImg(img, current_box)
                    cropPath = getCropImgPath(videoName, current_class)
                    cropped.save(cropPath)
                    print("in frame " + str(frame) + " find " + str(class_names[current_class]))
                    break
                elif((current_y - img_height * baseline_y) * (old_y - img_height * baseline_y) < 0):
                    
                    traffic_count[current_class] += 1
                    cropped = cropImg(img, current_box)
                    cropPath = getCropImgPath(videoName, current_class)
                    cropped.save(cropPath)
                    print("in frame " + str(frame) + " find " + str(class_names[current_class]))
                    break

# 获取所有class
def get_class():
    classes_path = os.path.expanduser('D:/yolo3-keras/model_data/coco_classes.txt')
    with open(classes_path) as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]
    return class_names

yolo = YOLO()

# video/1.avi 4辆
# video/2.mp4 17辆
# video/3.avi 

old_boxs = [] #前一帧box
old_classes = [] #前一帧class
baseline_x = 0.3 #x基准线
baseline_y = 0.7 #y基准线
img_width = 0 #图片宽度
img_height = 0 #图片高度
traffic_x_count = [0] * 80 #横向车道每一种class计数
traffic_y_count = [0] * 80 #纵向车道每一种class计数
traffic_count = [0] * 80
class_names = get_class()  #获取，每种物体的名称
videoName = ''

while True:
    video_path = input('Input videa filename:')
    try:
        num, img_width, img_height = videoToImg(video_path, 6)  #每6帧获取一张图片
    except:
        print('Open Error! Try again!')
        continue
    else:
        # traffic_x_count = [0] * 80 #横向车道每一种class计数
        # traffic_y_count = [0] * 80 #纵向车道每一种class计数 
        traffic_count = [0] * 80
        videoName = os.path.basename(video_path)
        inputFilePath = './output-img/' + videoName + '/'    #输入图片路径
        predictFilePath = "./predict-img/" + videoName + "/" #预测图片路径
        mkdir(predictFilePath)                               #创建预测图片路径
        object_rank = 0
        
        #创建json文件
        jsonadd = predictFilePath + videoName+'.json'        #存储所有物体的json文件的路径
        carplatej = predictFilePath + videoName+'car'+'.json'#存储识别出车牌的车辆的json文件路径
        jsonframe = predictFilePath + videoName+'jsonframe'+'.json'#存储每帧所含物体的json文件路径
        
        #创建临时json数组
        fuckjson = []                                        #存储所有物体的json文件的数组
        fuckjsoncar = []                                     #存储识别出车牌的车辆的json文件数组
        fuckframejson = []                                   #存储每帧所含物体的json文件数组
        
        #创建车牌数组
        tempcar = []
        for i in range(num):                                 #开始遍历每一帧

            imgpath = inputFilePath + str(i + 1) + '.jpg'    #输入文件路径
            image = Image.open(imgpath)                      #打开输入文件
            detected,object_rank = yolo.detect_image(image,predictFilePath,object_rank,
            fuckjson,
            fuckjsoncar,
            fuckframejson,
            tempcar
            )                                                #对图片进行检测
        
            # 对检测后图片的处理
            r_image = detected[0]
            mkdir(predictFilePath +'output/')
            
            outputPath = predictFilePath +'output/' + str(i + 1) + '.jpg'
            
            r_image.save(outputPath)                         #保存检测过后的图片

            # 检测车流量
            new_boxs = detected[1]                           #新的boxes
            new_classes = detected[2]                        #新的classes
            traffic(old_boxs, new_boxs, old_classes, new_classes, image, i)
            old_boxs = new_boxs                              #变成老的boxes
            old_classes = new_classes                        #变成老的classes
            #
            fuckframejson[i]["count"][0]=traffic_count[2]#车辆
            fuckframejson[i]["count"][1]=traffic_count[0]#行人
            fuckframejson[i]["count"][2]=traffic_count[3]#摩托车
            fuckframejson[i]["count"][3]=traffic_count[7]#卡车

            #imagetest = Image.open(outputPath) #显示主图像-测试
            
            #imagetest.show()

            '''
            fuckjson = []                                        #存储所有物体的json文件的数组
            fuckjsoncar = []                                     #存储识别出车牌的车辆的json文件数组
            fuckframejson = []                                   #存储每帧所含物体的json文件数组
            此时这一帧的所有信息都已更新到数组内部了

            '''


            
        #将json数组写入文件
        fr1 = open(jsonadd, 'w',encoding='utf-8') #打开json文件
        jsonstr = json.dumps(fuckjson)
        fr1.write(jsonstr)
        fr1.close()

        fr2 = open(carplatej, 'w',encoding='utf-8') #打开json文件
        jsonstr = json.dumps(fuckjsoncar)
        fr2.write(jsonstr)
        fr2.close()

        fr3 = open(jsonframe, 'w',encoding='utf-8') #打开json文件
        jsonstr = json.dumps(fuckframejson)
        fr3.write(jsonstr)
        fr3.close()
        for i in range(len(class_names)):          #遍历所有种类的数目，如果不等于零，则输出
            if traffic_count[i] != 0:
                print(class_names[i] + ': ' + str(traffic_count[i]))

yolo.close_session()


